package com.reebake.ideal.protect.web;

import cn.hutool.core.io.IoUtil;
import cn.hutool.core.util.StrUtil;
import com.reebake.ideal.crypto.entity.SecretKeyEntity;
import com.reebake.ideal.protect.core.ContentSecureService;
import com.reebake.ideal.servlet.core.ServletInputStreamWrapper;
import jakarta.servlet.ServletInputStream;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;

public class DecryptedRequestWrapper extends HttpServletRequestWrapper {
    private final ContentSecureService contentSecureService;
    private final SecretKeyEntity secretKeyEntity;
    private String plainBody;

    public DecryptedRequestWrapper(HttpServletRequest request, ContentSecureService contentSecureService, SecretKeyEntity secretKeyEntity) {
        super(request);
        this.contentSecureService = contentSecureService;
        this.secretKeyEntity = secretKeyEntity;
    }

    @Override
    public BufferedReader getReader() throws IOException {
        if (plainBody == null) {
            String contentType = super.getHeader(HttpHeaders.CONTENT_TYPE);
            String requestBody = readBody(contentType, super.getInputStream());
            plainBody = decryptRequestBody(requestBody);
        }
        return new BufferedReader(new InputStreamReader(getDecryptedServletInputStream()));
    }

    private String readBody(String contentType, InputStream is) {
        String requestBody = IoUtil.readUtf8(is);
        if(contentType.contains(MediaType.APPLICATION_JSON_VALUE)){
            requestBody = StrUtil.removePrefix(requestBody, "\"");
            requestBody = StrUtil.removeSuffix(requestBody, "\"");
        }
        return requestBody;
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        if (plainBody == null) {
            String contentType = super.getHeader(HttpHeaders.CONTENT_TYPE);
            String requestBody = readBody(contentType, super.getInputStream());
            plainBody = decryptRequestBody(requestBody);
        }
        return getDecryptedServletInputStream();
    }

    private String decryptRequestBody(String requestBody) throws IOException {
        try {
            return contentSecureService.decrypt(requestBody, secretKeyEntity);
        } catch (Exception e) {
            throw new IOException("Decryption failed", e);
        }
    }

    private ServletInputStream getDecryptedServletInputStream() throws IOException {
        return new ServletInputStreamWrapper(this.plainBody == null ? new byte[]{} : this.plainBody.getBytes(StandardCharsets.UTF_8));
    }
}
